import numpy as np
from torch.nn import functional as F
import torch
from model.fusions.fusions import MFB
import torch.nn as nn
#model中用到了
class ContextSelfAttn(nn.Module):#该模型接收一个表示视觉特征的张量 visual_feat，经过线性变换、注意力计算和汇总等操作产生一个表示视觉蒸馏结果的张量 v_distill。
    def __init__(self, module_dim=768):
        super(ContextSelfAttn, self).__init__()
        self.module_dim = module_dim
        self.v_proj = nn.Linear(module_dim, module_dim, bias=False)

        self.attn = nn.Linear(module_dim, 1)  #创建了线性层 attn 用于计算注意力分数
        self.activation = nn.ELU()
        self.dropout = nn.Dropout(0.15)

    def forward(self, visual_feat):
        """
        Args:
            question_rep: [Tensor] (batch_size, module_dim)
            visual_feat: [Tensor] (batch_size, num_of_clips, module_dim)
        return:
            visual_distill representation [Tensor] (batch_size, module_dim)
        """
        visual_feat = self.dropout(visual_feat) # TODO
        v_proj = self.v_proj(visual_feat)
        v_q_cat = self.activation(v_proj)

        attn = self.attn(v_q_cat)  # (bz, k, 1)#线性层 attn 对 v_q_cat 进行线性变换，得到注意力权重 attn。
        attn = F.softmax(attn, dim=1)  # (bz, k, 1)#使用 softmax 函数对 attn 进行归一化，得到概率分布。

        v_distill = (attn * visual_feat).sum(1)#将概率分布 attn 与视觉特征 visual_feat 相乘，并在第二维上求和，得到视觉蒸馏结果 v_distill。

        return v_distill   #模型利用自注意力机制，对每个视觉特征进行注意力计算，从而聚焦于重要的特征表示，并生成一个视觉蒸馏结果。

class SimpleOutputUnitOpenEnded(nn.Module): # concat->Linear classifier   用到了
    def __init__(self, module_dim=512, num_answers=1000):#num_answers 参数指定了分类器的输出类别数目
        super(SimpleOutputUnitOpenEnded, self).__init__()
       # 该模型接收一个表示问题嵌入的张量
       # question_embedding和一个表示视觉嵌入的张量visual_embedding，通过线性分类器将它们进行拼接并输出最终的结果张量out。
        self.question_proj = nn.Linear(module_dim, module_dim)

        self.classifier = nn.Sequential(nn.Dropout(0.15),
                                        nn.Linear(module_dim * 2, module_dim),
                                        nn.ELU(),
                                        nn.BatchNorm1d(module_dim),
                                        nn.Dropout(0.15),
                                        nn.Linear(module_dim, num_answers))
  #创建了一个包含多个层的顺序模型 classifier，其中包括一个 dropout 层、两个线性层、一个 ELU 激活函数、一个归一化层和一个 dropout 层。
    def forward(self, question_embedding, visual_embedding):
        question_embedding = self.question_proj(question_embedding)
        out = torch.cat([visual_embedding, question_embedding], 1)
        #使用 torch.cat 函数将视觉嵌入和问题嵌入在第二维上进行拼接，得到一个维度为 module_dim * 2 的张量 out。
        out = self.classifier(out)
   #将 out 输入到顺序模型 classifier 中，经过一系列的层操作后，得到最终的结果张量 out。
        return out


